import os
import sys
import json
import time
import subprocess
import requests
import zipfile
import io
from pathlib import Path
from typing import Optional, Dict, Any, List, Tuple
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from .logger import logger

class DriverManager:
    def __init__(self):
        self.driver_dir = Path("drivers")
        self.driver_dir.mkdir(exist_ok=True)
        self.chrome_driver_path = self.driver_dir / "chromedriver.exe"
        self._chrome_driver_path = None
        
        # ChromeDriver下载源配置
        self.endpoints = [
            "https://googlechromelabs.github.io/chrome-for-testing",
            "https://storage.googleapis.com/chrome-for-testing-public",
            "https://registry.npmmirror.com/-/binary/chromedriver",
            "https://cdn.npmmirror.com/binaries/chromedriver"
        ]
        
        # 配置请求会话
        self.session = requests.Session()
        retries = Retry(
            total=3,
            backoff_factor=0.5,
            status_forcelist=[500, 502, 503, 504]
        )
        self.session.mount('http://', HTTPAdapter(max_retries=retries))
        self.session.mount('https://', HTTPAdapter(max_retries=retries))

    def get_chrome_version(self) -> Optional[str]:
        """获取Chrome浏览器版本"""
        try:
            if sys.platform == "win32":
                import winreg
                try:
                    # 从注册表获取版本
                    key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, r"Software\Google\Chrome\BLBeacon")
                    version, _ = winreg.QueryValueEx(key, "version")
                    winreg.CloseKey(key)
                    logger.info(f"从注册表获取到Chrome版本: {version}")
                    return version
                except WindowsError:
                    # 从安装路径获取版本
                    paths = [
                        os.path.expandvars(r"%ProgramFiles%\Google\Chrome\Application\chrome.exe"),
                        os.path.expandvars(r"%ProgramFiles(x86)%\Google\Chrome\Application\chrome.exe"),
                        os.path.expandvars(r"%LocalAppData%\Google\Chrome\Application\chrome.exe")
                    ]
                    
                    for path in paths:
                        if os.path.exists(path):
                            try:
                                result = subprocess.check_output(
                                    ['powershell', '-command',
                                     f'(Get-Item "{path}").VersionInfo.FileVersion'],
                                    stderr=subprocess.STDOUT
                                ).decode().strip()
                                logger.info(f"从安装路径获取到Chrome版本: {result}")
                                return result
                            except subprocess.CalledProcessError:
                                continue
            
            logger.error("未能获取Chrome版本")
            return None
            
        except Exception as e:
            logger.error(f"获取Chrome版本失败: {str(e)}")
            return None

    def get_matching_driver_url(self, chrome_version: str) -> Optional[str]:
        """获取匹配的ChromeDriver下载URL"""
        try:
            chrome_major_version = chrome_version.split('.')[0]
            
            # 优先尝试淘宝镜像源
            taobao_endpoints = [
                "https://registry.npmmirror.com/-/binary/chromedriver",
                "https://cdn.npmmirror.com/binaries/chromedriver"
            ]
            
            for endpoint in taobao_endpoints:
                try:
                    # 获取版本列表
                    version_list_url = f"{endpoint}/"
                    response = self.session.get(version_list_url, timeout=10)
                    if response.status_code == 200:
                        versions_data = response.json()
                        # 查找匹配的版本
                        matching_version = None
                        # 处理淘宝镜像源的版本列表格式
                        if isinstance(versions_data, dict) and 'mirrors' in versions_data:
                            versions = [v['name'] for v in versions_data.get('mirrors', [])]
                        else:
                            versions = versions_data
                        
                        # 将版本号转换为元组以进行数值比较
                        def version_to_tuple(v):
                            try:
                                return tuple(map(int, v.split('.')))
                            except (AttributeError, ValueError):
                                return (0,)
                        
                        chrome_version_tuple = version_to_tuple(chrome_major_version)
                        matching_versions = []
                        
                        for version in versions:
                            if not isinstance(version, str):
                                continue
                            version_tuple = version_to_tuple(version)
                            if version_tuple[0] == chrome_version_tuple[0]:
                                matching_versions.append(version)
                        
                        # 选择最接近的版本
                        if matching_versions:
                            matching_version = max(matching_versions, key=version_to_tuple)
                        
                        if matching_version:
                            # 构建下载URL
                            download_url = f"{endpoint}/{matching_version}/chromedriver_win32.zip"
                            # 验证下载URL是否可用
                            head_response = self.session.head(download_url, timeout=10)
                            if head_response.status_code == 200:
                                return download_url
                except Exception as e:
                    logger.warning(f"尝试淘宝镜像源 {endpoint} 失败: {str(e)}")
                    continue
            
            # 如果淘宝镜像源失败，尝试其他端点
            for endpoint in self.endpoints:
                if endpoint in taobao_endpoints:
                    continue
                try:
                    known_good_versions_url = f"{endpoint}/known-good-versions-with-downloads.json"
                    response = self.session.get(known_good_versions_url, timeout=10)
                    if response.status_code == 200:
                        versions_data = response.json()
                        # 查找匹配的版本
                        for version in versions_data['versions']:
                            if version['version'].startswith(chrome_major_version + '.'):
                                # 获取Windows平台的下载URL
                                for download in version['downloads']['chromedriver']:
                                    if download['platform'] == 'win64':
                                        return download['url']
                except Exception as e:
                    logger.warning(f"尝试端点 {endpoint} 失败: {str(e)}")
                    continue
            
            raise Exception("所有下载源均无法访问或未找到匹配版本")
            
        except Exception as e:
            logger.error(f"获取ChromeDriver下载URL失败: {str(e)}")
            return None

    def download_chromedriver(self, chrome_version: str) -> bool:
        """下载ChromeDriver"""
        try:
            # 获取下载URL
            download_url = self.get_matching_driver_url(chrome_version)
            if not download_url:
                raise Exception("无法获取下载URL")
            
            # 下载ChromeDriver
            logger.info(f"开始下载ChromeDriver: {download_url}")
            max_retries = 3
            retry_delay = 2
            
            for attempt in range(max_retries):
                try:
                    response = self.session.get(download_url, timeout=30)
                    if response.status_code == 200:
                        break
                    raise Exception(f"下载失败，状态码: {response.status_code}")
                except Exception as e:
                    if attempt < max_retries - 1:
                        logger.warning(f"下载重试 {attempt + 1}/{max_retries}: {str(e)}")
                        time.sleep(retry_delay)
                        retry_delay *= 2
                        continue
                    raise
            
            # 解压ChromeDriver
            with zipfile.ZipFile(io.BytesIO(response.content)) as zip_file:
                chromedriver_path = None
                for file in zip_file.namelist():
                    if file.endswith('chromedriver.exe'):
                        chromedriver_path = file
                        break
                
                if not chromedriver_path:
                    raise Exception("ZIP文件中未找到chromedriver.exe")
                
                # 提取ChromeDriver
                zip_file.extract(chromedriver_path, self.driver_dir)
                
                # 移动到正确的位置
                if chromedriver_path != 'chromedriver.exe':
                    extracted_path = self.driver_dir / chromedriver_path
                    if self.chrome_driver_path.exists():
                        self.chrome_driver_path.unlink()
                    extracted_path.rename(self.chrome_driver_path)
            
            logger.info(f"ChromeDriver下载成功: {self.chrome_driver_path}")
            return True
            
        except Exception as e:
            logger.error(f"下载ChromeDriver失败: {str(e)}")
            return False

    def get_driver_path(self) -> Optional[str]:
        """获取ChromeDriver路径，如果不存在则延迟下载"""
        try:
            # 如果驱动存在，直接返回路径
            if self.chrome_driver_path.exists():
                return str(self.chrome_driver_path)
            
            # 驱动不存在时，返回None，让调用者决定是否需要下载
            return None
            
        except Exception as e:
            logger.error(f"获取ChromeDriver路径失败: {str(e)}")
            return None

    def ensure_driver(self) -> Optional[str]:
        """确保ChromeDriver可用，必要时下载"""
        try:
            # 检查现有驱动
            driver_path = self.get_driver_path()
            if driver_path:
                return driver_path
            
            # 获取Chrome版本并下载驱动
            chrome_version = self.get_chrome_version()
            if not chrome_version:
                raise Exception("无法获取Chrome版本")
                
            if not self.download_chromedriver(chrome_version):
                raise Exception("下载ChromeDriver失败")
            
            return str(self.chrome_driver_path)
            
        except Exception as e:
            logger.error(f"确保ChromeDriver可用失败: {str(e)}")
            return None

# 创建全局驱动管理器实例
driver_manager = DriverManager()

# 导出便捷函数
def get_chrome_driver_path() -> Optional[str]:
    return driver_manager.get_driver_path()

def get_chrome_version() -> Optional[str]:
    return driver_manager.get_chrome_version()
